import akshare as ak
import pandas as pd
import requests
import re
import time
from retrying import retry
import os
from datetime import datetime, timedelta
import numpy as np

def display_welcome():
    """显示欢迎信息"""
    print("=" * 60)
    print("          股票历史市盈率查询工具")
    print("=" * 60)
    print("支持输入格式：")
    print("  - 6位数字代码：600519（贵州茅台）")
    print("  - 带市场前缀：sh600519 或 sz000001")
    print("  - 输入 'quit' 或 'exit' 退出程序")
    print("=" * 60)
    print("功能：获取股票历史数据并计算动态市盈率")

def validate_stock_code(code):
    """验证股票代码格式"""
    code = code.strip().lower()
    
    if code in ['quit', 'exit', 'q']:
        return 'exit'
    
    # 处理带市场前缀的代码
    if code.startswith(('sh', 'sz')):
        market = code[:2]
        number = code[2:]
        if number.isdigit() and len(number) == 6:
            return f"{market}{number}"
    
    # 处理纯数字代码
    elif code.isdigit() and len(code) == 6:
        if code.startswith(('6', '9')):
            return f"sh{code}"
        else:
            return f"sz{code}"
    
    return None

@retry(stop_max_attempt_number=3, wait_fixed=2000)
def get_stock_data_akshare():
    """使用akshare获取股票数据，带重试机制"""
    try:
        return ak.stock_zh_a_spot_em()
    except Exception as e:
        print(f"  akshare获取数据失败: {e}")
        raise e

def get_stock_name_from_akshare(stock_code, stock_data):
    """从akshare数据中获取股票名称"""
    try:
        pure_code = stock_code[2:]
        stock_info = stock_data[stock_data['代码'] == pure_code]
        
        if not stock_info.empty and '名称' in stock_info.columns:
            return stock_info['名称'].values[0]
        return "未知股票"
        
    except Exception as e:
        print(f"  获取股票名称出错: {e}")
        return "未知股票"

@retry(stop_max_attempt_number=3, wait_fixed=2000)
def get_historical_price_data(stock_code, start_date="20200101", end_date=None):
    """获取历史价格数据"""
    try:
        pure_code = stock_code[2:]
        if end_date is None:
            end_date = datetime.now().strftime("%Y%m%d")
        
        print(f"  获取历史价格数据 {start_date} 至 {end_date}...")
        
        # 使用akshare获取历史价格数据
        df = ak.stock_zh_a_hist(symbol=pure_code, period="daily", 
                               start_date=start_date, end_date=end_date, adjust="qfq")
        
        if df is not None and not df.empty:
            df['日期'] = pd.to_datetime(df['日期'])
            df.set_index('日期', inplace=True)
            print(f"  成功获取 {len(df)} 条历史价格数据")
            return df
        else:
            return None
            
    except Exception as e:
        print(f"  获取历史价格数据失败: {e}")
        return None

@retry(stop_max_attempt_number=3, wait_fixed=2000)
def get_financial_data(stock_code):
    """获取财务数据"""
    try:
        pure_code = stock_code[2:]
        print(f"  获取财务数据...")
        
        # 获取最新财务指标
        df = ak.stock_financial_analysis_indicator(symbol=pure_code)
        
        if df is not None and not df.empty:
            # 获取最新的净利润数据
            latest_data = df.iloc[0]
            print(f"  成功获取财务数据")
            return latest_data
        return None
            
    except Exception as e:
        print(f"  获取财务数据失败: {e}")
        return None

def get_total_shares(stock_code):
    """获取总股本（更准确的方法）"""
    try:
        pure_code = stock_code[2:]
        
        # 方法1: 使用stock_individual_info_em接口
        stock_info = ak.stock_individual_info_em(symbol=pure_code)
        if stock_info is not None and not stock_info.empty:
            for _, row in stock_info.iterrows():
                if '总股本' in str(row['item']) or 'total shares' in str(row['item']).lower():
                    value_str = str(row['value'])
                    if '亿' in value_str:
                        total_shares = float(value_str.replace('亿', '')) * 1e8
                        print(f"  获取总股本: {total_shares/1e8:.2f} 亿股")
                        return total_shares
        
        # 方法2: 使用实时行情数据估算
        print("  尝试从实时数据估算总股本...")
        stock_spot = ak.stock_zh_a_spot_em()
        stock_data = stock_spot[stock_spot['代码'] == pure_code]
        if not stock_data.empty and '总市值' in stock_data.columns and '最新价' in stock_data.columns:
            market_cap = stock_data['总市值'].values[0] * 1e8  # 转换为元
            price = stock_data['最新价'].values[0]
            if price > 0:
                total_shares = market_cap / price
                print(f"  估算总股本: {total_shares/1e8:.2f} 亿股")
                return total_shares
        
        # 方法3: 使用默认值（根据股票类型）
        if pure_code.startswith('6'):  # 沪市
            default_shares = 12.56e8  # 约12.56亿股（贵州茅台的大致股本）
        else:
            default_shares = 10e8  # 默认10亿股
        
        print(f"  使用默认总股本: {default_shares/1e8:.2f} 亿股")
        return default_shares
        
    except Exception as e:
        print(f"  获取总股本失败: {e}")
        return 12.56e8  # 贵州茅台的近似股本

def get_net_profit(stock_code):
    """获取净利润数据（更准确的方法）"""
    try:
        pure_code = stock_code[2:]
        
        # 方法1: 获取最新财务报告的净利润
        financial_data = get_financial_data(stock_code)
        if financial_data is not None and '净利润' in financial_data:
            net_profit = financial_data['净利润'] * 1e8  # 转换为元
            print(f"  获取净利润: {net_profit/1e8:.2f} 亿元")
            return net_profit
        
        # 方法2: 根据不同股票提供相对准确的估算值
        if pure_code == '600519':  # 贵州茅台
            net_profit = 700e8  # 约700亿元（2023年数据）
        elif pure_code == '000001':  # 平安银行
            net_profit = 400e8  # 约400亿元
        elif pure_code == '600036':  # 招商银行
            net_profit = 1300e8  # 约1300亿元
        else:
            # 对于其他股票，使用相对合理的估算
            net_profit = 50e8  # 默认50亿元
        
        print(f"  使用估算净利润: {net_profit/1e8:.2f} 亿元")
        return net_profit
        
    except Exception as e:
        print(f"  获取净利润失败: {e}")
        return 700e8  # 默认值

def calculate_pe_ratio(close_price, total_shares, net_profit):
    """计算市盈率"""
    if net_profit > 0 and total_shares > 0:
        market_cap = close_price * total_shares
        pe_ratio = market_cap / net_profit
        return pe_ratio
    return None

def calculate_historical_pe(stock_code, price_df):
    """计算历史市盈率"""
    if price_df is None or price_df.empty:
        return None
    
    # 获取总股本和净利润
    total_shares = get_total_shares(stock_code)
    net_profit = get_net_profit(stock_code)
    
    print(f"  计算参数: 总股本={total_shares/1e8:.2f}亿股, 净利润={net_profit/1e8:.2f}亿元")
    
    # 复制数据框
    result_df = price_df.copy()
    
    # 计算市盈率
    pe_ratios = []
    market_caps = []
    
    for close_price in result_df['收盘']:
        pe_ratio = calculate_pe_ratio(close_price, total_shares, net_profit)
        market_cap = close_price * total_shares if close_price > 0 else 0
        
        pe_ratios.append(pe_ratio)
        market_caps.append(market_cap / 1e8)  # 转换为亿元
    
    result_df['动态市盈率'] = pe_ratios
    result_df['总市值(亿元)'] = market_caps
    result_df['净利润(亿元)'] = net_profit / 1e8
    
    # 移除异常的市盈率值
    result_df = result_df[result_df['动态市盈率'] > 0]
    result_df = result_df[result_df['动态市盈率'] < 1000]  # 移除异常大的值
    
    print(f"  成功计算 {len(pe_ratios)} 个市盈率数据点")
    return result_df

def save_historical_data_to_file(df, stock_code, stock_name, filename=None):
    """将历史数据保存到文件"""
    if df is None or df.empty:
        print("  无数据可保存")
        return False
    
    if filename is None:
        # 生成默认文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{stock_code}_{stock_name}_历史市盈率_{timestamp}.csv"
    
    try:
        # 确保输出目录存在
        output_dir = "stock_pe_history"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        filepath = os.path.join(output_dir, filename)
        
        # 复制DataFrame以避免修改原始数据
        save_df = df.copy()
        
        # 添加日期列
        save_df['日期'] = save_df.index
        
        # 选择要保存的列
        columns_to_save = ['日期', '收盘', '开盘', '最高', '最低', '成交量', '成交额', '动态市盈率', '总市值(亿元)', '净利润(亿元)']
        existing_columns = [col for col in columns_to_save if col in save_df.columns]
        
        # 保存数据
        save_df[existing_columns].to_csv(filepath, encoding='utf-8-sig', index=False)
        print(f"  数据已保存到: {filepath}")
        print(f"  共保存 {len(save_df)} 条记录")
        print(f"  时间范围: {save_df['日期'].min()} 至 {save_df['日期'].max()}")
        
        return True
        
    except Exception as e:
        print(f"  保存文件失败: {e}")
        return False

def display_data_statistics(df, stock_name):
    """显示数据统计信息"""
    if df is None or df.empty:
        return
    
    print(f"\n📊 {stock_name} 数据统计:")
    print(f"   数据期间: {df.index.min()} 至 {df.index.max()}")
    print(f"   记录数量: {len(df)} 条")
    
    if '动态市盈率' in df.columns:
        pe_data = df['动态市盈率'].dropna()
        if len(pe_data) > 0:
            print(f"   动态市盈率 - 最新: {pe_data.iloc[-1]:.2f}")
            print(f"   动态市盈率 - 平均: {pe_data.mean():.2f}")
            print(f"   动态市盈率 - 最高: {pe_data.max():.2f}")
            print(f"   动态市盈率 - 最低: {pe_data.min():.2f}")
            
            # 市盈率分析
            latest_pe = pe_data.iloc[-1]
            if latest_pe < 0:
                print("   💡 分析: 市盈率为负，公司可能亏损")
            elif latest_pe < 15:
                print("   💡 分析: 市盈率较低，可能被低估")
            elif latest_pe < 25:
                print("   💡 分析: 市盈率适中")
            elif latest_pe < 40:
                print("   💡 分析: 市盈率较高，注意风险")
            else:
                print("   💡 分析: 市盈率很高，投资需谨慎")
    
    if '收盘' in df.columns:
        close_data = df['收盘'].dropna()
        if len(close_data) > 0:
            print(f"   收盘价 - 最新: {close_data.iloc[-1]:.2f}")
            print(f"   收盘价 - 最高: {close_data.max():.2f}")
            print(f"   收盘价 - 最低: {close_data.min():.2f}")
    
    if '总市值(亿元)' in df.columns:
        market_cap = df['总市值(亿元)'].dropna()
        if len(market_cap) > 0:
            print(f"   总市值 - 最新: {market_cap.iloc[-1]:.2f} 亿元")

def get_time_period():
    """获取用户选择的时间周期"""
    print("\n请选择时间周期:")
    print("1. 最近半年")
    print("2. 最近1年")
    print("3. 最近2年")
    print("4. 最近3年")
    
    choice = input("请选择(1-4, 默认2): ").strip()
    
    today = datetime.now()
    if choice == "1":
        start_date = (today - timedelta(days=180)).strftime("%Y%m%d")
    elif choice == "3":
        start_date = (today - timedelta(days=365*2)).strftime("%Y%m%d")
    elif choice == "4":
        start_date = (today - timedelta(days=365*3)).strftime("%Y%m%d")
    else:
        start_date = (today - timedelta(days=365)).strftime("%Y%m%d")
    
    return start_date

def main():
    """主函数"""
    display_welcome()
    
    # 预先加载一次股票数据
    print("正在初始化数据...")
    try:
        stock_data = get_stock_data_akshare()
        print("数据加载成功！")
    except Exception as e:
        print(f"数据加载失败: {e}")
        stock_data = None
    
    while True:
        try:
            # 获取用户输入
            user_input = input("\n请输入股票代码: ").strip()
            
            # 检查退出命令
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("感谢使用，再见！")
                break
            
            # 验证股票代码
            validated_code = validate_stock_code(user_input)
            
            if validated_code == 'exit':
                print("感谢使用，再见！")
                break
            elif validated_code is None:
                print("❌ 股票代码格式错误！请输入6位数字代码或带市场前缀的代码")
                continue
            
            # 获取股票名称
            stock_name = "未知股票"
            if stock_data is not None:
                stock_name = get_stock_name_from_akshare(validated_code, stock_data)
            
            print(f"正在处理: {stock_name} ({validated_code})")
            
            # 获取时间周期
            start_date = get_time_period()
            
            # 获取历史价格数据
            print(f"正在获取历史价格数据...")
            price_data = get_historical_price_data(validated_code, start_date=start_date)
            
            if price_data is not None and not price_data.empty:
                # 计算市盈率
                print(f"正在计算动态市盈率...")
                pe_data = calculate_historical_pe(validated_code, price_data)
                
                if pe_data is not None and not pe_data.empty:
                    # 显示统计信息
                    display_data_statistics(pe_data, stock_name)
                    
                    # 保存到文件
                    save_success = save_historical_data_to_file(pe_data, validated_code, stock_name)
                    
                    if save_success:
                        print("✅ 历史数据获取和计算完成！")
                        
                        # 显示文件内容预览
                        try:
                            output_dir = "stock_pe_history"
                            files = os.listdir(output_dir)
                            if files:
                                latest_file = max([os.path.join(output_dir, f) for f in files], key=os.path.getctime)
                                preview_df = pd.read_csv(latest_file, nrows=5)
                                print(f"\n📋 文件预览（前5行）:")
                                print(preview_df.to_string(index=False))
                        except:
                            pass
                else:
                    print("❌ 市盈率计算失败")
            else:
                print("❌ 无法获取该股票的历史价格数据")
            
            # 询问是否继续
            continue_input = input("\n是否继续查询其他股票？(y/n): ").strip().lower()
            if continue_input not in ['y', 'yes', '是', '']:
                print("感谢使用，再见！")
                break
                
        except KeyboardInterrupt:
            print("\n\n程序被用户中断，再见！")
            break
        except Exception as e:
            print(f"发生错误: {e}")
            continue

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"程序运行出错: {e}")
